import os
import numpy as np
import cv2

input_dir = "/disk527/sdb1/a804_cbf/datasets/Lunar_Crater_Detection_Data/LRO_DATA/train"

output_dir = "/disk527/sdb1/a804_cbf/datasets/haworth/train"
resolution = 200
if not os.path.exists(output_dir):
    os.makedirs(output_dir)
for id, file in enumerate(os.listdir(input_dir)):
    idx = 0
    file = file.strip().split(".")[0]
    img = cv2.imread(os.path.join(input_dir, file + ".png"), cv2.IMREAD_GRAYSCALE)
    for i in range(img.shape[0] // resolution):
        for j in range(img.shape[1] // resolution):
            patch = img[
                i * resolution : (i + 1) * resolution,
                j * resolution : (j + 1) * resolution,
            ]
            cv2.imwrite(os.path.join(output_dir, f"{id}_{i}_{j}.png"), patch)
    with open(os.path.join(input_dir, file + ".txt"), "r") as f:
        f.readline()
        for line in f:
            line = line.strip().split(",")
            x1 = int(float(line[1]))
            y1 = int(float(line[2]))
            x2 = int(float(line[3]) + x1)
            y2 = int(float(line[4]) + y1)
            i = x1 // resolution
            j = y1 // resolution
            ii = x2 // resolution
            jj = y2 // resolution
            if x2 - x1 > 30 or y2 - y1 > 30:
                continue
            if ii == i and jj == j:
                with open(os.path.join(output_dir, f"{id}_{i}_{j}.txt"), "a") as writer:
                    print(
                        f"{x1 - i * resolution} {y1 - j * resolution} {x2 - x1} {y2 - y1}",
                        file=writer,
                        flush=True,
                    )
            elif ii > 0 and jj > 0 and i > 0 and j > 0:
                with open(
                    os.path.join(output_dir, f"{id}_{ii}_{jj}.txt"), "a"
                ) as writer:
                    print(
                        f"{x1 - ii * resolution} {y1 - jj * resolution} {x2 - x1} {y2 - y1}",
                        file=writer,
                        flush=True,
                    )
                with open(os.path.join(output_dir, f"{id}_{i}_{j}.txt"), "a") as writer:
                    print(
                        f"{x1 - i * resolution} {y1 - j * resolution} {x2 - x1} {y2 - y1}",
                        file=writer,
                        flush=True,
                    )
